import os
import sys
import time

import matplotlib.pyplot as plt
import numpy as np
from torchsummary import summary

root_path = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.append(root_path)
import torch
from sklearn.metrics import f1_score, accuracy_score, precision_score, recall_score, classification_report
from torch import nn
from torch.optim import AdamW
from torch.utils.data import Dataset, DataLoader
from tqdm import tqdm
from transformers import BertTokenizer, BertModel, BertConfig
from utils import id2label, label2id, ensure_pdir_exist
from config import conf
from bert_yt.b_config import b_conf


def load_data(path):
    data = []
    with open(path, 'r', encoding='utf-8') as f:
        for line in tqdm(f, f'load data from {path}'):
            line = line.strip()
            if not line:
                continue
            split = line.split('\t')
            data.append((split[0], split[1]))
    print(f'数据长度：{len(data)}')
    return data


class MyDataset(Dataset):
    def __init__(self, path):
        self.data = load_data(path)
        self.len = len(self.data)

    def __len__(self):
        return self.len

    def __getitem__(self, index):
        x = self.data[index][1]
        y = label2id[self.data[index][0]]
        return x, y


def _test_dataset():
    for i in MyDataset(conf.val_path):
        print(i)
        break


device = conf.device
# device = 'cpu'
print(f'b_conf.pretrained_model_path={b_conf.pretrained_model_path}')
bert_config = BertConfig(b_conf.pretrained_model_path)
bert_tokenizer = BertTokenizer.from_pretrained(b_conf.pretrained_model_path)
bert_model = BertModel.from_pretrained(b_conf.pretrained_model_path)


def collate_fn(batch):
    texts = [x[0] for x in batch]
    labels = [x[1] for x in batch]
    tokens = bert_tokenizer.batch_encode_plus(
        texts,
        return_tensors='pt',
        max_length=b_conf.max_length,
        padding=True,
        truncation=True,
    )
    return tokens['input_ids'], tokens['attention_mask'], torch.LongTensor(labels)


def get_dataloader(dataset, batch_size=b_conf.batch_size, num_workers=4):
    return DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=True,
        collate_fn=collate_fn,
        num_workers=num_workers,
        pin_memory=True
    )


class BertModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.bert = bert_model
        self.fc = nn.Linear(bert_config.hidden_size, len(id2label))

    def forward(self, input_ids, attention_mask, return_hidden=False):
        with torch.no_grad():
            outputs = self.bert(input_ids, attention_mask)
        cls_hidden = outputs.last_hidden_state[:, 0]
        logits = self.fc(cls_hidden)
        if return_hidden:
            return logits, cls_hidden
        return logits


def _test_model():
    dataset = MyDataset(conf.val_path)
    dataloader = get_dataloader(dataset)
    model = BertModel()
    for input_ids, attention_mask, labels in dataloader:
        print(input_ids, attention_mask, labels)
        pred = model(input_ids, attention_mask)
        print(pred)
        break


def validata_on_dev(model, val_dataloader):
    pred_list = []
    label_list = []
    for input_ids, attention_mask, labels in val_dataloader:
        input_ids, attention_mask, labels = input_ids.to(device), attention_mask.to(device), labels.to(device)
        with torch.no_grad():
            model.eval()
            pred = model(input_ids, attention_mask)
        pred_list += pred.argmax(dim=1).tolist()
        label_list += labels.tolist()
    f1 = f1_score(label_list, pred_list, average='micro')
    print(f'f1 on dev = {f1:.4f}')
    return f1


def train():
    dataloader = get_dataloader(MyDataset(conf.train_path))
    val_dataloader = get_dataloader(MyDataset(conf.val_path))
    model = BertModel().to(device)
    optimizer = AdamW(model.parameters(), lr=b_conf.lr)
    criterion = nn.CrossEntropyLoss()

    for p in bert_model.parameters():
        p.requires_grad_(False)

    epoch_dev_max_f1 = 0.0
    train_f1 = []
    dev_f1 = []
    x_label = []
    for i in range(b_conf.epochs):
        total_loss = 0.0
        total_item = 0
        pred_list = []
        label_list = []
        batch_dev_final_f1 = 0.0
        for j, (input_ids, attention_mask, labels) in tqdm(enumerate(dataloader), f'iter dataloader'):
            input_ids, attention_mask, labels = input_ids.to(device), attention_mask.to(device), labels.to(device)
            model.train()
            pred = model(input_ids, attention_mask)
            loss = criterion(pred, labels)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            total_loss += loss * len(input_ids)
            total_item += len(input_ids)
            pred_list += pred.argmax(dim=1).tolist()
            label_list += labels.tolist()
            if (j + 1) % b_conf.log_batch == 0:
                print(f'epoch: {i + 1}, iter: {j + 1}, avg_loss: {total_loss / total_item}')
                f1 = f1_score(label_list, pred_list, average='micro')
                print(f'f1 on train = {f1:.4f}')
                batch_dev_final_f1 = validata_on_dev(model, val_dataloader)
                train_f1.append(f1)
                dev_f1.append(batch_dev_final_f1)
                x_label.append(f'{i + 1}-{j + 1}')

        pretty_draw(train_f1, dev_f1, x_label, b_conf.f1_img_path[:-4] + f'{i + 1}.png')
        if batch_dev_final_f1 > epoch_dev_max_f1:
            epoch_dev_max_f1 = batch_dev_final_f1
            ensure_pdir_exist(b_conf.model_path)
            torch.save(model.state_dict(), b_conf.model_path)
            print(f'更新验证集最大f1={epoch_dev_max_f1}，保存模型')

    # epoch4 更新验证集最大f1=0.9518，保存模型
    pass


_model = None


def get_model():
    global _model
    if _model is None:
        _model = BertModel().to(device)
        _model.load_state_dict(torch.load(b_conf.model_path, map_location=conf.device_str))
    return _model


def evaluate(model=get_model(), return_incorrect=0):
    model.eval()
    dataset = MyDataset(conf.test_path)
    batch_size = 2048
    dataloader = get_dataloader(dataset, batch_size=batch_size)

    pred_list = []
    label_list = []
    incorrect = []
    start = time.time()
    with torch.no_grad():
        for i, (input_ids, attention_mask, labels) in enumerate(tqdm(dataloader, 'test...')):
            input_ids, attention_mask = input_ids.to(device, non_blocking=True), attention_mask.to(device,
                                                                                                   non_blocking=True)
            pred = model(input_ids, attention_mask)
            pred_list += pred.argmax(dim=1).tolist()
            label_list += labels.tolist()
            if return_incorrect > 0 and len(incorrect) < return_incorrect:
                for j, (yp, yt) in enumerate(zip(pred_list, label_list)):
                    if yp != yt:
                        text_idx = i * b_conf.batch_size + j
                        incorrect.append((text_idx, yp, yt))
            if len(incorrect) > return_incorrect:
                print('错误用例数量达标提前退出')
                break
    cost = time.time() - start
    print(f'batch_size={batch_size}，每秒处理数={len(pred_list) / cost}')
    acc_score = accuracy_score(label_list, pred_list)
    print(f'准确率：{acc_score}')
    pre_score = precision_score(label_list, pred_list, average='micro')
    print(f'精确率：{pre_score}')
    rec_score = recall_score(label_list, pred_list, average='micro')
    print(f'召回率：{rec_score}')
    f1 = f1_score(label_list, pred_list, average='micro')
    print(f'F1-score：{f1}')
    target_names = [id2label[i] for i in range(len(id2label))]
    c_report = classification_report(label_list, pred_list, target_names=target_names)
    print(f'分类报告：\n{c_report}')
    """
准确率：0.9573
精确率：0.9573
召回率：0.9573
F1-score：0.9573
分类报告：
              precision    recall  f1-score   support

          体育       1.00      1.00      1.00      1000
          娱乐       0.99      0.99      0.99      1000
          家居       0.96      0.89      0.92      1000
          房产       0.90      0.88      0.89      1000
          教育       0.96      0.93      0.95      1000
          时尚       0.98      0.97      0.98      1000
          时政       0.92      0.97      0.94      1000
          游戏       0.99      0.98      0.98      1000
          科技       0.94      0.99      0.96      1000
          财经       0.94      0.98      0.96      1000

    accuracy                           0.96     10000
   macro avg       0.96      0.96      0.96     10000
weighted avg       0.96      0.96      0.96     10000
    """
    return dataset, incorrect[:return_incorrect]


def pretty_draw(train_f1, dev_f1, x_label, path=b_conf.f1_img_path, x_label_desc='epoch-batch'):
    plt.figure(figsize=(10, 6), dpi=100)
    plt.plot(x_label, train_f1,
             color='#1f77b4',  # 蓝色
             linestyle='-',
             linewidth=2,
             marker='o',
             markersize=8,
             label='Train F1')

    plt.plot(x_label, dev_f1,
             color='#ff7f0e',  # 橙色
             linestyle='--',
             linewidth=2,
             marker='s',
             markersize=8,
             label='Dev F1')

    plt.title('F1 Score Trend During Training', fontsize=14, pad=20)
    plt.xlabel(x_label_desc, fontsize=12)
    plt.ylabel('F1 Score', fontsize=12)
    plt.xticks(x_label)  # 显示所有 epoch
    plt.ylim(0, 1)  # F1 范围 0-1
    plt.grid(True, linestyle='--', alpha=0.7)
    plt.legend(loc='lower right', fontsize=10)
    plt.tight_layout()
    ensure_pdir_exist(path)
    plt.savefig(path)
    plt.close()


def predict(text):
    model = get_model()
    input_ids, attention_mask, _ = collate_fn([(text, -1)])
    input_ids, attention_mask = input_ids.to(device), attention_mask.to(device)
    with torch.no_grad():
        return id2label[model(input_ids, attention_mask).argmax(dim=1).item()]


def _test_predict():
    for text, y_true in zip(conf.test_texts, conf.test_results):
        result = predict(text.strip())
        print(f'predict text={text[:100]}\nresult={result}\tcorrect={result == y_true}')
    """
predict text=
            内墙涂料抽检五成不合格 警惕苯含量超标昨日，广东省工商局公布了省内销售的涂料、胶粘剂的最新抽检结果。结果显示，内墙涂料不合格率最高，为50%，而白乳胶不合格率最低，为25%，
result=家居     correct=True
predict text=
    打女李宇春：从功夫狗熊变身功夫熊猫(组图)"春春在《十月围城》里吃了两次饭，流了一次眼泪，剩下的时间都是在打群架。"虽然前几天《十月围城》南京首映庆典上李宇春有事没能参加，不过剧组显然很重视
result=娱乐     correct=True
predict text=
    行业政策与经济周期的博弈21世纪经济报道 徐广蓉 根据全球最著名的房地产周期研究专家哈里森的研究，一个完整的房地产周期是18年左右：房价先上涨7年，然后发生短期的下跌，再经历5年的快速上涨，
result=房产     correct=True
    """


def _test_cr():
    # 设置随机种子以便结果可复现
    np.random.seed(42)
    y_true = np.random.randint(0, 10, size=30)

    # 生成30个样本的预测标签（模拟模型预测，包含一些错误）
    y_pred = y_true.copy()
    # 随机选择5个样本，将其预测结果修改为错误类别
    error_indices = np.random.choice(30, size=5, replace=False)
    for i in error_indices:
        # 确保生成的错误标签与真实标签不同
        wrong_labels = [j for j in range(10) if j != y_true[i]]
        y_pred[i] = np.random.choice(wrong_labels)

    # 定义10个标签的名称列表
    target_names = [f"类别{i}" for i in range(10)]

    # 生成并打印分类报告
    report = classification_report(y_true, y_pred, target_names=target_names)
    print(report)


def ana_incorrect():
    dataset, incorrect = evaluate(return_incorrect=20)
    for text_idx, yp, yt in incorrect:
        text = dataset[text_idx][0]
        print('=' * 50)
        print(f'预测错误用例 text={text[:100]}\ny_pred={id2label[yp]}\t y_true={id2label[yt]}')

    """
==================================================
预测错误用例 text=骑士老板狂笑捋詹姆斯 丹别在我车库里撒野!(图)新浪体育讯北京时间3月30日消息  热火客场以90比102爆冷输给骑士，错过了超越凯尔特人登上东部第二的良机。赛后，各方人士通过微博第一时间发表了看法。
y_pred=房产      y_true=财经
==================================================
预测错误用例 text=美媒评现役假摔之王 骑士影帝榜首火箭两人入前五新浪体育讯北京时间3月31日，如果调查NBA的假摔王，相信每个球迷心中都会有那么几个名字。近日，美国某著名篮球媒体针对所有现役球员，展开了一项匿名调查：谁
y_pred=教育      y_true=房产
==================================================
预测错误用例 text=专访帕特森：和姚相提并论很骄傲 要全力赢下每场新浪体育讯北京时间4月3日(休斯敦时间4月2日消息)，帕特里克-帕特森最近表现出色，之前一度还代替受伤的斯科拉担任先发。他接受了新浪体育的专访，对于最近火
y_pred=家居      y_true=时尚
==================================================
预测错误用例 text=李：火箭别无选择只能拼 胜马刺靠的是季后赛强度新浪体育讯北京时间4月3号(休斯顿当地时间4月2号消息)，周六休斯顿火箭队在丰田中心进行了两个小时的例行训练，在结束了一个客场之旅之后，火箭队又将回到一天
y_pred=房产      y_true=家居
==================================================
预测错误用例 text=查尔默斯米勒火线复出 斯帅：皇帝每晚能打48分钟新浪体育讯北京时间4月4日，在客场与篮网的比赛前，热火主帅埃里克-斯波尔斯特拉表示，两位伤病迈克-米勒和马里奥-查尔默斯都将进入出场名单。“迈克和马里奥
y_pred=游戏      y_true=科技
==================================================
预测错误用例 text=直击-洛瑞堵枪眼被撞飞挂彩 嘴唇缝6针后重返赛场新浪体育讯北京时间4月4日消息，火箭继续主场作战，迎战东部第五的老鹰。在赛季即将结束的时候，为了季后赛理论上的可能，火箭仍在全力的一拼。帕楚里亚抢下篮板
y_pred=家居      y_true=房产
==================================================
预测错误用例 text=国王vs火箭前瞻：休城冲八不容有失 马丁欲破魔咒新浪体育讯北京时间4月5日，休斯顿火箭队明天将在主场丰田中心迎战萨克拉门托国王队。火箭队依然在季后赛的边缘徘徊，而国王在西部排名倒数第二，战绩仅仅优于森
y_pred=房产      y_true=财经
==================================================
预测错误用例 text=步行者为锁定季后赛欢庆 02届探花：不枉我等了9年新浪体育讯正所谓“春风得意马蹄疾，一朝看尽长安花”。在今天以136-112大胜华盛顿奇才后，由于夏洛特山猫加时以102-111负于奥兰多魔术，步行者锁
y_pred=财经      y_true=房产
==================================================
预测错误用例 text=湖人一败正式无缘西部第一 马刺坐收渔利锁定王位新浪体育讯北京时间4月7日，又见冷门！洛杉矶湖人在客场87-95不敌金州勇士，遭遇三连败。全明星后一度17胜1负的卫冕冠军突然掉了链子，背靠背输给了西部倒
y_pred=科技      y_true=家居
==================================================
预测错误用例 text=现场：布莱切9分5板勇猛依旧 桑帅咆哮裁判引技犯 新浪体育讯北京时间4月10日，华盛顿奇才回到主场迎战亚特兰大老鹰。两支球队都是背靠背作战，此前的三次交手中，老鹰取得3胜0负。开场跳球后，奇才率先得球
y_pred=房产      y_true=财经
==================================================
预测错误用例 text=雷霆vs湖人首发：科比PK杜兰特 洛城双塔迎来强敌新浪体育讯北京时间4月11日，洛杉矶湖人坐镇主场，迎战上赛季季后赛首轮对手俄克拉荷马雷霆。赛前，两队公布了各自首发阵容，这次昔日对手再碰面都是尽遣主力
y_pred=教育      y_true=房产
==================================================
预测错误用例 text=堂堂国家队主帅爆粗口 辽宁媒体质问他在侮辱谁？本报上海电 记者李倩倩报道 辽沪焦点战，场上队员打得激烈，场下教练也挺较劲。上半场，卢卡斯推倒边强，就一度引得两位主帅振臂呼喊义愤填膺。而后一次双方犯规的
y_pred=家居      y_true=时尚
==================================================
预测错误用例 text=雷霆五连胜终止名定第四 小牛后劲十足胜黄蜂32分新浪体育讯北京时间4月14日，雷霆在常规赛时间最后2分钟丧失9分的优势，最终加时不敌雄鹿，排名定为西部第四；小牛客场击败黄蜂，暂列西部第二，黄蜂则掉到第
y_pred=房产      y_true=家居
==================================================
预测错误用例 text=留不留火箭都得去努力面对 姚明：生活就像白开水驻休斯敦记者郭海舟4月13日报道 “生活就像白开水，哪儿有那么多刺激的，只能自己接受，自己找点乐趣呗，”此前一场比赛前，记者们和姚明聊到篮球、工作之外的生
y_pred=游戏      y_true=科技
==================================================
预测错误用例 text=传球失误也彰显大心脏 詹姆斯浪费他超想象力空接新浪体育讯迈阿密热火97-89战胜费城76人，迈克-毕比拿下8分。尽管毕比防守能力颇让人担心，热火主帅埃里克-斯波尔斯特拉依然选择让他担任热火季后赛首发控
y_pred=家居      y_true=房产
==================================================
预测错误用例 text=膝盖狂抽积水巴恩斯仍请战：只要能动我就能上场!新浪体育讯北京时间4月17日，据ESPN报道，卫冕冠军洛杉矶湖人队将在明天开始季后赛的征程，不过板凳上的悍将马特-巴恩斯面临着严重的膝伤问题。在本赛季常规
y_pred=房产      y_true=财经
==================================================
预测错误用例 text=父亲女友见证里程碑之战 最硬一招他只手左右胜负新浪体育讯相比于上一场的末节狂飙，德克-诺维茨基今天从一开场就不断的尝试加热自己的手感。于是，我们看见了一个在关键的第二节与第四节顺利开启天王模式的诺维茨
y_pred=财经      y_true=房产
==================================================
预测错误用例 text=科比：为奥多姆骄傲 鱼总豪言要让冠军为喇嘛增色新浪体育讯北京时间4月20日(洛杉矶时间4月19日)消息，奥多姆获得最佳第六人，湖人队友和教练菲尔-杰克逊都对此发表看法。杰克逊：我很高兴，他值得获得这个
y_pred=科技      y_true=家居
==================================================
预测错误用例 text=罗伊为抱怨上场时间事件道歉 主帅不计前嫌挺爱徒新浪体育讯北京时间4月22日消息，与小牛的前两场季后赛，开拓者昔日老大布兰顿-罗伊一共只得到了34分钟上场时间。而已经被逐渐边缘化的罗伊在昨日接受当地媒体
y_pred=房产      y_true=娱乐
==================================================
预测错误用例 text=佩顿：湖人将是本赛季总冠军 一对一我能防住罗斯新浪体育讯北京时间4月22日，美国NBA传奇球星加里-佩顿继续自己的中国之旅，继19日作为嘉宾参加上海车展东风风神H30 Cross上市仪式之后，今天来到
y_pred=房产      y_true=家居
    """
    pass


def _test_num_workers():
    dataset = MyDataset(conf.train_path)
    import os
    print(f"逻辑 CPU 数量: {os.cpu_count()}")
    for i in range(0, 41, 4):
        print(f'测试 num_workers: {i}')
        dataloader = get_dataloader(dataset, num_workers=i)
        for j, (input_ids, attention_mask, labels) in tqdm(enumerate(dataloader), f'iter dataloader'):
            input_ids, attention_mask = input_ids.to(device, non_blocking=True), attention_mask.to(device,
                                                                                                   non_blocking=True)
    """
测试 num_workers: 0
iter dataloader: 49it [03:03,  3.74s/it]
测试 num_workers: 4
iter dataloader: 49it [00:47,  1.02it/s]
测试 num_workers: 8
iter dataloader: 49it [00:28,  1.72it/s]
测试 num_workers: 12
iter dataloader: 49it [00:18,  2.58it/s]
测试 num_workers: 16
iter dataloader: 49it [00:16,  3.02it/s]
测试 num_workers: 20
iter dataloader: 49it [00:14,  3.33it/s]
测试 num_workers: 24
iter dataloader: 49it [00:17,  2.74it/s]
测试 num_workers: 28
iter dataloader: 49it [00:18,  2.59it/s]
测试 num_workers: 32
iter dataloader: 49it [00:20,  2.40it/s]
测试 num_workers: 36
iter dataloader: 49it [00:21,  2.29it/s]
测试 num_workers: 40
iter dataloader: 49it [00:23,  2.07it/s]

    """


if __name__ == '__main__':
    # _test_dataset()
    # _test_model()
    # _test_num_workers()
    # train()
    evaluate()
    # _test_predict()
    # _test_cr()
    # ana_incorrect()
    pass
